import pygame
import sys

import food as c_food
import snake as c_snake
import dqn as c_dqn


def game_over():
    pygame.quit()
    sys.exit()


def game_start():
    window_title = 'Snake-AI'
    window_size = (500, 500)
    grid_num = (10, 10)                  # 64列，48行
    grid_size = (50, 50)

    pygame.init()
    pg_clock = pygame.time.Clock()

    main_window = pygame.display.set_mode(window_size)
    pygame.display.set_caption(window_title)

    FOOD = c_food.food()
    SNAKE = c_snake.snake()
    DQN = c_dqn.dqn()

    move_dir = 0                                # 0,1,2,3    上下左右
    _move_dir = 0
    whether_eat = 0                             # 0未吃，1吃了
    whether_die = 0                             # 0活着，1die

    body_list, not_body_list = SNAKE.init_body(grid_num)
    whether_eat, food_list = FOOD.make_food(1, not_body_list, [])

    s = SNAKE.get_s(body_list, food_list, grid_num)
    r = 0
    d_ = grid_num[0]+grid_num[1]

    while True:

        move_dir = DQN.choose_action(s)

        if (_move_dir in [0, 1] and move_dir in [0, 1]) or (_move_dir in [2, 3] and move_dir in [2, 3]):
            continue

        not_body_list, body_list, whether_eat, whether_die = SNAKE.move_step(not_body_list, body_list, food_list, move_dir, grid_num)

        r, d_ = SNAKE.get_r(whether_eat, whether_die, r, body_list, food_list, grid_num, d_)
        s_ = SNAKE.get_s(body_list, food_list, grid_num)

        DQN.store_ds(s, move_dir, r, s_)
        if DQN.ds_st_learn == 1:
            DQN.learn()
        s = s_
        _move_dir = move_dir

        if whether_die == 1:
            body_list, not_body_list = SNAKE.init_body(grid_num)
            whether_eat, food_list = FOOD.make_food(1, not_body_list, [])
            whether_eat, whether_die, r, s, d_ = 0, 0, 0, SNAKE.get_s(body_list, food_list, grid_num), grid_num[0]+grid_num[1]
        whether_eat, food_list = FOOD.make_food(whether_eat, not_body_list, food_list)

        main_window.fill((0, 0, 0))
        FOOD.draw_food(food_list, main_window, grid_size)
        SNAKE.draw_body(body_list, main_window, grid_size)
        pygame.display.update()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()

        if len(body_list) >= 4 and DQN.learn_step >= 1000:
            pg_clock.tick(15)


if __name__ == '__main__':
    game_start()